import getopt
import sys
from datetime import datetime, timedelta

import pandas as pd
from dateutil.relativedelta import relativedelta

from air_web.config.config import config
from air_web.data_platform import init_db, sql_engine

orgno_typeid_15min_columns = """
  `org_no` int(11) NOT NULL COMMENT '区域编码',
  `org_name` varchar(255) DEFAULT NULL COMMENT '区域名称',
  `ad_org_name` varchar(255) DEFAULT NULL COMMENT '行政区域名称',
  `p_org_no` int(11) DEFAULT NULL COMMENT '父级区域编码',
  `type_id` int(11) NOT NULL COMMENT '行业类型编码',
  `type_code` varchar(64) DEFAULT NULL COMMENT '行业类型',
  `p_type_id` int(11) DEFAULT NULL COMMENT '父级行业类型编码',
  `data_time` datetime NOT NULL COMMENT '时间',
  `time_display_name` varchar(64) NOT NULL COMMENT '展示时间',
  `p_total_sum` double DEFAULT NULL COMMENT '总负荷总和',
  `p_kt_sum` double DEFAULT NULL COMMENT '空调负荷总和',
  `p_base_sum` double DEFAULT NULL COMMENT '其他负荷总和',
  `p_total_count` float DEFAULT NULL COMMENT '总负荷数据条数',
  `p_kt_count` float DEFAULT NULL COMMENT '空调负荷数据条数',
  `p_base_count` float DEFAULT NULL COMMENT '其他负荷数据条数',
  `is_day_max` int(11) NOT NULL COMMENT '是否当天最大',
  PRIMARY KEY (`org_no`,`type_id`,`data_time`),
  KEY `idx_otm_df` (`data_time`)
"""

baseline_orgno_typeid_15min_columns = """
  `org_no` int(11) NOT NULL COMMENT '区域编码',
  `org_name` varchar(255) DEFAULT NULL COMMENT '区域名称',
  `ad_org_name` varchar(255) DEFAULT NULL COMMENT '行政区域名称',
  `p_org_no` int(11) DEFAULT NULL COMMENT '父级区域编码',
  `type_id` int(11) NOT NULL COMMENT '行业类型编码',
  `type_code` varchar(64) DEFAULT NULL COMMENT '行业类型',
  `p_type_id` int(11) DEFAULT NULL COMMENT '父级行业类型编码',
  `data_time` datetime NOT NULL COMMENT '时间',
  `time_display_name` varchar(64) NOT NULL COMMENT '展示时间',
  `p_total_sum` double DEFAULT NULL COMMENT '总负荷总和',
  `p_kt_sum` double DEFAULT NULL COMMENT '空调负荷总和',
  `p_total_count` float DEFAULT NULL COMMENT '总负荷数据条数',
  `p_kt_count` float DEFAULT NULL COMMENT '空调负荷数据条数',
  `is_day_max` int(11) NOT NULL COMMENT '是否当天最大',
  PRIMARY KEY (`org_no`,`type_id`,`data_time`),
  KEY `idx_otm_df` (`data_time`)
"""


spec_type_15min_columns = """
  `org_no` int(11) NOT NULL COMMENT '区域编码',
  `org_name` varchar(255) DEFAULT NULL COMMENT '区域名称',
  `p_org_no` int(11) DEFAULT NULL COMMENT '父级区域编码',
  `type_id` int(11) NOT NULL COMMENT '行业类型编码',
  `type_code` varchar(64) DEFAULT NULL COMMENT '行业类型',
  `data_time` datetime NOT NULL COMMENT '时间',
  `p_total_sum` double DEFAULT NULL COMMENT '总负荷总和',
  `p_kt_sum` double DEFAULT NULL COMMENT '空调负荷总和',
  `p_base_sum` double DEFAULT NULL COMMENT '其他负荷总和',
  PRIMARY KEY (`org_no`,`type_id`,`data_time`),
  KEY `idx_otm_df` (`data_time`)
"""

spec_cons_15min_columns = """
  `org_no` int(11) NOT NULL COMMENT '区域编码',
  `org_name` varchar(255) DEFAULT NULL COMMENT '区域名称',
  `p_org_no` int(11) DEFAULT NULL COMMENT '父级区域编码',
  `type_id` int(11) NOT NULL COMMENT '行业类型编码',
  `type_code` varchar(64) DEFAULT NULL COMMENT '行业类型',
  `measure_code` int(11) NOT NULL COMMENT '调控措施编码',
  `measure_name` varchar(64) DEFAULT NULL COMMENT '调控措施名称',
  `data_time` datetime NOT NULL COMMENT '时间',
  `data_date` date NOT NULL COMMENT '日期',
  `p_total_sum` double DEFAULT NULL COMMENT '总负荷总和',
  `p_kt_sum` double DEFAULT NULL COMMENT '空调负荷总和',
  `p_total_baseline` double DEFAULT NULL COMMENT '总负荷基线总和',
  `p_kt_baseline` double DEFAULT NULL COMMENT '空调负荷基线总和',
  `is_day_max` int(11) NOT NULL COMMENT '是否当天最大',
  PRIMARY KEY (`org_no`,`type_id`,`measure_code`,`data_time`),
  KEY `idx_otm_df` (`data_date`)
"""

powerload_virtual_columns = """
  `data_time` datetime(0) NOT NULL,
  `p_total` double NULL DEFAULT NULL,
  `p_kt` double NULL DEFAULT NULL,
  `p_base` double NULL DEFAULT NULL,
  `p_std_left` float NULL DEFAULT NULL,
  `p_std_right` float NULL DEFAULT NULL,
  `kt_ratio` float NULL DEFAULT NULL,
  `cons_no` varchar(64) NOT NULL,
  `on5` int(11) NULL DEFAULT NULL,
  `on7` int(11) NULL DEFAULT NULL,
  `type_id` int(11) NULL DEFAULT NULL,
  `is_day_max` int(11) NULL DEFAULT NULL,
  `org_no` int(11) NULL DEFAULT NULL,
  `outlier` int(11) NULL DEFAULT NULL,
   PRIMARY KEY (`cons_no`, `data_time`)
"""

powerload_predict_virtual_columns = """
  `data_time` datetime(0) NOT NULL,
  `p_total` double NULL DEFAULT NULL,
  `p_kt` double NULL DEFAULT NULL,
  `cons_no` varchar(64) NOT NULL,
  `on5` int(11) NULL DEFAULT NULL,
  `on7` int(11) NULL DEFAULT NULL,
  `type_id` int(11) NULL DEFAULT NULL,
  `is_day_max` int(11) NULL DEFAULT NULL,
   PRIMARY KEY (`cons_no`, `data_time`)
"""


cons_ymd_stat_columns = """
  `cons_no` varchar(64) NOT NULL COMMENT '用户编码',
  `on5` int(11) DEFAULT NULL COMMENT '地市编码',
  `on7` int(11) DEFAULT NULL COMMENT '区县编码',
  `date_type` int(11) NOT NULL COMMENT '日期类型',
  `data_date` date NOT NULL COMMENT '日期',
  `max_p_kt` float NULL DEFAULT NULL COMMENT '最大空调负荷值',
  `max_p_kt_total` float NULL DEFAULT NULL COMMENT '最大空调负荷对应的总负荷',
  `max_p_kt_time` datetime(0) NOT NULL COMMENT '最大空调负荷时间点',
  `max_p_kt_bl_total` float NULL DEFAULT NULL COMMENT '最大空调负荷对应的总负荷基线',
  `max_p_kt_bl_kt` float NULL DEFAULT NULL COMMENT '最大空调负荷对应的空调负荷基线',
  `max_p_total` float NULL DEFAULT NULL COMMENT '最大总负荷值',
  `max_p_total_kt` float NULL DEFAULT NULL COMMENT '最大总负荷值对应的空调负荷',
  `max_p_total_time` datetime(0) NOT NULL COMMENT '最大总负荷时间点',
  `max_p_total_bl_total` float NULL DEFAULT NULL COMMENT '最大总负荷对应的总负荷基线',
  `max_p_total_bl_kt` float NULL DEFAULT NULL COMMENT '最大总负荷对应的空调负荷基线',
  `max_p_total_rank` int default 999999 COMMENT '最大总负荷排序',
  `avg_p_kt` float NULL DEFAULT NULL COMMENT '平均空调负荷',
  `avg_p_total` float NULL DEFAULT NULL COMMENT '平均总负荷',
  `min_p_kt` float NULL DEFAULT NULL COMMENT '最小空调负荷值',
  `min_p_kt_total` float NULL DEFAULT NULL COMMENT '最小空调负荷值对应的总负荷',
  `min_p_kt_time` datetime(0) NULL DEFAULT NULL COMMENT '最小空调负荷时间点',
  `min_p_total` float NULL DEFAULT NULL COMMENT '最小总负荷值',
  `min_p_total_kt` float NULL DEFAULT NULL COMMENT '最小总负荷值对应的空调负荷',
  `min_p_total_time` datetime(0) NULL DEFAULT NULL COMMENT '最小总负荷时间点',
  `max_tmp` float NULL DEFAULT NULL COMMENT '最高温度',
  `min_tmp` float NULL DEFAULT NULL COMMENT '最低温度',
  `is_base_day` int(11) NULL DEFAULT NULL COMMENT '是否是基准日:0空调日,1基准日,-1异常日',
  `outlier` int(11) NULL DEFAULT NULL COMMENT '是否是异常点:0正常,1异常',
  PRIMARY KEY (`cons_no`, `date_type`, `data_date`),
  KEY `idx_on5` (`on5`),
  KEY `idx_on7` (`on7`)
"""


p_96_list = [f"`p_{str(i+1)}` float NULL DEFAULT NULL" for i in range(96)]

cons_res_power_96_columns = f"""
  {','.join(p_96_list)},
  `data_date` date NOT NULL COMMENT '数据日期',
  `data_type` int(11) NOT NULL COMMENT '数据类型：1p_total,2p_kt,3p_base,4p_std_left,5p_std_right,6p_kt_est,7p_base_est',
  `cons_no` varchar(64) NOT NULL COMMENT '户号',
  `on5` int(11) NOT NULL COMMENT '地市编码',
  `on7` int(11) NOT NULL COMMENT '区县编码',
  `org_no` int(11) NULL DEFAULT NULL COMMENT '供电所编码',
  `type_id` int(11) NOT NULL COMMENT '小行业编码',
   PRIMARY KEY (`cons_no`, `data_date`, `data_type`)
"""

cons_res_baseline_columns = f"""
  {','.join(p_96_list)},
  `data_date` date NOT NULL COMMENT '数据日期',
  `data_type` int(11) NOT NULL COMMENT '数据类型：1p_total,2p_kt',
  `cons_no` varchar(64) NOT NULL COMMENT '户号',
  `on5` int(11) NOT NULL COMMENT '地市编码',
  `on7` int(11) NOT NULL COMMENT '区县编码',
  `org_no` int(11) NULL DEFAULT NULL COMMENT '供电所编码',
  `type_id` int(11) NOT NULL COMMENT '小行业编码',
   PRIMARY KEY (`cons_no`, `data_date`, `data_type`)
"""

cons_res_predict_columns = f"""
  {','.join(p_96_list)},
  `data_date` date NOT NULL COMMENT '数据日期',
  `data_type` int(11) NOT NULL COMMENT '数据类型：1p_total,2p_kt',
  `cons_no` varchar(64) NOT NULL COMMENT '户号',
  `on5` int(11) NOT NULL COMMENT '地市编码',
  `on7` int(11) NOT NULL COMMENT '区县编码',
  `org_no` int(11) NULL DEFAULT NULL COMMENT '供电所编码',
  `type_id` int(11) NOT NULL COMMENT '小行业编码',
   PRIMARY KEY (`cons_no`, `data_date`, `data_type`)
"""


def get_add_partition_sql(start_date, end_date, freq="1M"):
    date_list = pd.date_range(
        start=start_date, end=end_date, freq=freq
    ).strftime("%Y-%m-%d")
    add_partition = f""
    for date in date_list:
        p_date = date[0:10] if freq == "1d" else date[0:7]
        partition_name = f"p{p_date.replace('-', '')}"
        partition_date = datetime.strptime(date, "%Y-%m-%d") + timedelta(
            days=1
        )
        add_partition += f"PARTITION {partition_name} VALUES LESS THAN (TO_DAYS('{partition_date}')), \n"
    add_partition = add_partition[:-3]  # 去掉最后一个逗号和空格
    return add_partition


def gen_org_table(
    table_name: str,
    start_date: str,
    end_date: str,
    partition_column: str,
    columns: str,
    freq="1M",
    sql_engine=sql_engine,
):
    # 创建主表
    create_table = f"""CREATE TABLE if not exists {table_name} (
    {columns}
    ) PARTITION BY RANGE (TO_DAYS({partition_column})) ("""

    # 动态创建分区
    add_partition = get_add_partition_sql(start_date, end_date, freq)
    create_sql = f"{create_table}\n{add_partition}\n);"
    print(create_sql)
    sql_engine.engine.execute(create_sql)


def add_partition(
    table_name: str,
    start_date: str,
    end_date: str,
    freq="1M",
    sql_engine=sql_engine,
):
    # 动态创建分区
    add_partition = get_add_partition_sql(start_date, end_date, freq)
    sql = f"ALTER TABLE {table_name} ADD PARTITION (\n {add_partition}\n);"
    print(sql)
    sql_engine.engine.execute(sql)


class GenTablePartion:
    def __init__(self, *args, **kwargs):
        pass

    def main(
        self, start_date="2019-01-01", end_date="2025-01-01", is_add=False
    ):
        cons_sql_engine = init_db(config["MYSQL_CONS_DB_NAME"])
        orgno_typeid_15min_table = config.get(
            "ORGNO_TYPEID_15MIN", "orgno_typeid_15min"
        )
        baseline_orgno_typeid_15min_table = config.get(
            "BASELINE_ORGNO_TYPEID_15MIN", "baseline_orgno_typeid_15min"
        )
        orderly_power_15min = config.get(
            "ORDERLY_POWER_15MIN", "orderly_power_15min"
        )
        demand_power_15min = config.get(
            "DEMAND_POWER_15MIN", "demand_power_15min"
        )
        virtual_table = config.get(
            "POWERLOAD_VIRTUAL_TABLE", "aclr_res_power_virtual"
        )
        virtual_predict_table = config.get(
            "POWERLOAD_PREDICT_VIRTUAL_TABLE", "aclr_res_power_predict_virtual"
        )
        cons_ymd_stat_table = config.get("CONS_YMD_STAT", "cons_ymd_stat")
        spec_type_15min_table = config.get(
            "SPEC_TYPE_15MIN", "spec_type_15min"
        )
        cons_res_power_96_table = "aclr_res_power_96"
        cons_res_baseline_table = "aclr_res_baseline"
        cons_res_predict_table = "aclr_res_predict_96"

        partition_column = "data_time"

        if is_add:
            add_partition(orgno_typeid_15min_table, start_date, end_date)
            add_partition(
                baseline_orgno_typeid_15min_table, start_date, end_date
            )
            add_partition(orderly_power_15min, start_date, end_date)
            add_partition(demand_power_15min, start_date, end_date)
            add_partition(virtual_table, start_date, end_date)
            add_partition(virtual_predict_table, start_date, end_date)
            add_partition(cons_ymd_stat_table, start_date, end_date, "1d")
            add_partition(spec_type_15min_table, start_date, end_date)
            # 用户详情
            add_partition(
                cons_res_power_96_table,
                start_date,
                end_date,
                "1d",
                cons_sql_engine,
            )
            add_partition(
                cons_res_baseline_table,
                start_date,
                end_date,
                "1d",
                cons_sql_engine,
            )
            add_partition(
                cons_res_predict_table,
                start_date,
                end_date,
                "1d",
                cons_sql_engine,
            )
        else:
            gen_org_table(
                orgno_typeid_15min_table,
                start_date,
                end_date,
                partition_column,
                orgno_typeid_15min_columns,
            )
            gen_org_table(
                baseline_orgno_typeid_15min_table,
                start_date,
                end_date,
                partition_column,
                baseline_orgno_typeid_15min_columns,
            )
            gen_org_table(
                orderly_power_15min,
                start_date,
                end_date,
                partition_column,
                spec_cons_15min_columns,
            )
            gen_org_table(
                demand_power_15min,
                start_date,
                end_date,
                partition_column,
                spec_cons_15min_columns,
            )
            gen_org_table(
                virtual_table,
                start_date,
                end_date,
                partition_column,
                powerload_virtual_columns,
            )
            gen_org_table(
                virtual_predict_table,
                start_date,
                end_date,
                partition_column,
                powerload_predict_virtual_columns,
            )
            gen_org_table(
                cons_ymd_stat_table,
                start_date,
                end_date,
                "data_date",
                cons_ymd_stat_columns,
                "1d",
            )
            gen_org_table(
                spec_type_15min_table,
                start_date,
                end_date,
                partition_column,
                spec_type_15min_columns,
            )
            # 用户详情
            gen_org_table(
                cons_res_power_96_table,
                start_date,
                end_date,
                "data_date",
                cons_res_power_96_columns,
                "1d",
                sql_engine=cons_sql_engine,
            )
            gen_org_table(
                cons_res_baseline_table,
                start_date,
                end_date,
                "data_date",
                cons_res_baseline_columns,
                "1d",
                sql_engine=cons_sql_engine,
            )
            gen_org_table(
                cons_res_predict_table,
                start_date,
                end_date,
                "data_date",
                cons_res_predict_columns,
                "1d",
                sql_engine=cons_sql_engine,
            )


if __name__ == "__main__":
    start_date = "2019-01-01"  # 开始日期
    end_date = "2025-01-01"  # 结束日期
    is_add = False  # 为False是第一次建表，为True则是建表以后添加新的分区

    opts, args = getopt.getopt(sys.argv[1:], "s:e:a")
    for opt, val in opts:
        if opt == "-s":
            start_date = val
        elif opt == "-e":
            end_date = val
        elif opt == "-a":
            is_add = True

    GenTablePartion().main(start_date, end_date, is_add)
